# !/usr/bin/python
# -*- coding: utf-8 -*-

import os
import shutil
import warnings

import fitz
import win32api
import win32con
from cnocr import CnOcr
from cnocr.utils import read_img


# ---把pdf每页都转换成一张图片--- #
def pdf_to_img(input_pdf_name, output_file):
    pdf = fitz.open(input_pdf_name + '.pdf')
    page = pdf.page_count  # 读取pdf的页数
    for i in range(0, page):
        page_fn = pdf.load_page(i)  # 读取第i页
        pix = page_fn.get_pixmap()
        pix.save(os.path.join(output_file, str(i) + '.png'))  # 用函数链接输出路径，避免不同系统的分隔符差异带来的报错


# ---识别图片中的文字并且return--- #
def img_ocr(input_pic):
    ocr = CnOcr()  # 可以填入参数context="cpu/gpu/cuda:0"等，用别的资源训练，默认为cpu。
    img_fp = input_pic  # 输入图片路径
    img = read_img(img_fp)  # 读取图片
    res = ocr.ocr(img)  # res是list（列表）
    len_res = len(res)  # 文本的行数
    text = ""
    for i in range(0, len_res):
        line_tuple = res[i]  # line_tuple是tuple（元组）
        s = ""
        single_line = s.join(line_tuple[0])
        text += single_line
    return text  # text为识别出来的文本，此处return可以给下面写入txt


# ---将识别出来的文字写入txt--- #
def write_to_txt(input_pdf_name, text):
    with open(input_pdf_name + '.txt', 'a', encoding='utf-8') as file:  # 用a模式是续写而不是w(重新写入),编码为utf-8
        file.write(text)
        file.write('\n')  # 写完一张图片里面的内容就换行


# ---获取当前路径下所有pdf文件--- #
def get_pdf_list():
    pdf_list = [x for x in os.listdir('.') if os.path.isfile(x) and os.path.splitext(x)[1] == '.pdf']
    return pdf_list


# ---忽略报错--- #
def ignore_warning():
    warnings.filterwarnings("ignore")


# ---连接路径--- #
def join_path(path1, path2):
    joined_path = os.path.join(path1, path2)
    return joined_path


# ---创建目录--- #
def mk_dir(filename):
    os.mkdir(filename)


# ---删除目录--- #
def rm_dir(filename):
    shutil.rmtree(filename)


# ---弹窗提示完成--- #
def finish():
    win32api.MessageBox(0, "已完成PDF转化", "已完成", win32con.MB_ICONASTERISK)


def start_ocr():
    ignore_warning()
    pdf_list = get_pdf_list()
    for i in range(0, len(pdf_list)):
        name = os.path.splitext(pdf_list[i])[0]
        input_pdf_name = join_path(os.path.curdir, name)
        output_file = join_path(os.path.curdir, 'pdf_2_image')
        input_file = output_file
        mk_dir(output_file)
        pdf_to_img(input_pdf_name, output_file)
        pic = os.listdir(input_file)
        for pic_fn in pic:
            text = img_ocr(join_path(input_file, pic_fn))
            write_to_txt(input_pdf_name, text)
        rm_dir(output_file)
    finish()


if __name__ == "__main__":
    start_ocr()
